# coding:utf-8
from torch.utils.data import DataLoader  # 导入 PyTorch 的 DataLoader，用于批量加载数据
from transformers import default_data_collator  # 导入 Hugging Face 的默认数据整理器
from data_handle.data_preprocess import *  # 导入自定义的数据预处理模块
from pet_config import *  # 导入项目配置文件

pc = ProjectConfig()  # 实例化项目配置类，获取配置参数
tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)  # 加载预训练的 Tokenizer

def get_data():
    """
    获取并处理训练和验证数据集，返回 DataLoader。
    """
    prompt = open(pc.prompt_file, 'r', encoding='utf8').readlines()[0].strip()  # 读取 Prompt 文件的第一行内容
    hard_template = HardTemplate(prompt=prompt)  # 创建硬模板对象，用于模板化处理
    dataset = load_dataset('text', data_files={'train': pc.train_path, 'dev': pc.dev_path})  # 加载训练和验证数据集
    # print(dataset)
    # print(f'Prompt is -> {prompt}')  # 打印 Prompt 内容（被注释掉）

    # 定义偏函数，固定部分参数
    new_func = partial(convert_example,
                       tokenizer=tokenizer,  # 使用的 Tokenizer
                       hard_template=hard_template,  # 使用的硬模板
                       max_seq_len=pc.max_seq_len,  # 最大序列长度
                       max_label_len=pc.max_label_len)  # 最大标签长度

    # 对数据集进行批量转换处理
    dataset = dataset.map(new_func, batched=True)

    train_dataset = dataset["train"]  # 获取训练数据集
    dev_dataset = dataset["dev"]  # 获取验证数据集
    # print('train_dataset', train_dataset[:2])  # 打印前两个样本（被注释掉）
    # print('*'*80)

    # 创建训练数据的 DataLoader
    train_dataloader = DataLoader(
        train_dataset,  # 训练数据集
        shuffle=True,  # 是否打乱数据
        collate_fn=default_data_collator,  # 数据整理器，将数据转换为 Tensor 类型
        batch_size=pc.batch_size  # 批量大小
    )
    # 创建验证数据的 DataLoader
    dev_dataloader = DataLoader(
        dev_dataset,  # 验证数据集
        collate_fn=default_data_collator,  # 数据整理器
        batch_size=pc.batch_size  # 批量大小
    )
    return train_dataloader, dev_dataloader  # 返回训练和验证的 DataLoader


if __name__ == '__main__':
    train_dataloader, dev_dataloader = get_data()  # 调用函数获取 DataLoaders
    print(len(train_dataloader))  # 打印训练 DataLoader 的批次数量
    print(len(dev_dataloader))  # 打印验证 DataLoader 的批次数量
    for i, value in enumerate(train_dataloader):  # 遍历训练 DataLoader
        print(i)  # 打印当前批次索引
        print(value)  # 打印当前批次数据
        print(value['input_ids'].dtype)  # 打印 input_ids 的数据类型
        break  # 只打印第一个批次后退出循环
# 实现了一个完整的数据加载和预处理流程，用于深度学习模型的训练和验证。它通过读取 Prompt 模板和数据集，结合 Hugging Face 的工具链（如 AutoTokenizer 和 default_data_collator），将原始文本数据转换为模型所需的输入格式，并通过 DataLoader 提供批量加载功能。整个流程支持动态配置（如批量大小、最大序列长度等），具有良好的扩展性和通用性，适用于基于 Prompt 的自然语言处理任务（如文本分类、填空预测等）。